# writed by zealous guys.
import xml.etree.ElementTree as ET
import numpy as np
from pycocotools.coco import COCO
from tqdm import tqdm
import torch
import glob
from collections import defaultdict
import json

inf = 1000000


def id2name(coco):
    classes = dict()
    classes_id = []
    for cls in coco.dataset["categories"]:
        classes[cls["id"]] = cls["name"]

    for key in classes.keys():
        classes_id.append(key)
    return classes, classes_id


def load_dataset(path, types="voc"):
    dataset = []
    if types == "voc":
        for xml_file in glob.glob("{} /*xml".format(path)):
            tree = ET.parse(xml_file)
            # 图片高度
            height = int(tree.findtext("./size/height"))
            # 图片宽度
            width = int(tree.findtext("./size/width"))

            for obj in tree.iter("object"):
                # 偏移量
                xmin = int(obj.findtext("bndbox/xmin")) / width
                ymin = int(obj.findtext("bdbox/ymin")) / height
                xmax = int(obj.findtext("bndbox/xmax")) / width
                ymax = int(obj.findtext("bndbox/ymax")) / height
                xmin = np.float64(xmin)
                ymin = np.float64(ymin)
                xmax = np.float64(xmax)
                ymax = np.float64(ymax)
                if xmax == xmin or ymax == ymin:
                    print(xml_file)
                # 将Anchor的长宽放入dateset，运行kmeans获得Anchor
                dataset.append([xmax - xmin, ymax - ymin])

    if types == "coco":

        if not isinstance(path, list):
            coco = COCO(path)

            classes, classes_id = id2name(coco)
            print(classes)
            print("class_ids:", classes_id)

            img_ids = coco.getImgIds()
            print(len(img_ids))

            for imgId in img_ids:
                i = 0
                img = coco.loadImgs(imgId)[i]
                height = img["height"]
                width = img["width"]
                i = i + 1
                if imgId % 500 == 0:
                    print("process {} images".format(imgId))
                annIds = coco.getAnnIds(imgIds=img["id"], iscrowd=None)
                anns = coco.loadAnns(annIds)
                # time.sleep(0.2)
                for ann in anns:
                    if "bbox" in ann:
                        bbox = ann["bbox"]
                        """
                         coco:
                        annotation: [x, y, width, height]
                        """
                        ann_width = bbox[2]
                        ann_height = bbox[3]
                        if ann_width == 0 or ann_height == 0:
                            continue
                        # 偏移量
                        ann_width = np.float64(ann_width / width)
                        ann_height = np.float64(ann_height / height)
                        dataset.append([ann_width, ann_height])
                    else:
                        raise ValueError("coco no bbox -- wrong!!!")

        if isinstance(path, list):
            for x in range(len(path) - 1):
                x = 1
                coco = COCO(path[x])

                classes, classes_id = id2name(coco)
                print(classes)
                print("class_ids:", classes_id)

                img_ids = coco.getImgIds()
                print(len(img_ids))

                for imgId in img_ids:
                    i = 0
                    img = coco.loadImgs(imgId)[i]
                    # print(i)
                    # print(img)
                    height = img["height"]
                    width = img["width"]
                    # i = i + 1
                    if imgId % 500 == 0:
                        print("process {} images".format(imgId))
                    annIds = coco.getAnnIds(imgIds=img["id"], iscrowd=None)
                    anns = coco.loadAnns(annIds)
                    # time.sleep(0.2)
                    for ann in anns:
                        if "bbox" in ann:
                            if True:  # 'landm' in ann:
                                bbox = ann["bbox"]

                                # landm=ann['landm']
                                """
                                 coco:
                                annotation: [x, y, width, height]
                                """
                                if True:  # landm[0]>0:
                                    ann_width = bbox[2]
                                    ann_height = bbox[3]
                                    if ann_width <= 0 or ann_height <= 0:
                                        continue
                                    # if ann_width<0 or ann_height<0:
                                    #    print(bbox)
                                    #    print("1")
                                    #    exit()
                                    # 偏移量
                                    ann_width = np.float64(ann_width)  # / width)
                                    ann_height = np.float64(ann_height)  # / height)
                                    dataset.append([ann_width, ann_height])
                            else:
                                bbox = ann["bbox"]
                                """
                                 coco:
                                annotation: [x, y, width, height]
                                """
                                ann_width = bbox[2]
                                ann_height = bbox[3]
                                if ann_width <= 0 or ann_height <= 0:
                                    continue
                                # if ann_width<0 or ann_height<0:
                                #    print(bbox)
                                #    print("1")
                                #    exit()
                                # 偏移量
                                ann_width = np.float64(ann_width / width)
                                ann_height = np.float64(ann_height / height)
                                dataset.append([ann_width, ann_height])

                        else:
                            raise ValueError("coco no bbox -- wrong!!!")

    return np.array(dataset)


def split_dataset(path, types, split_flag, reverse):
    """
    Args:
        path: 待分割的数据路径
        types: 数据的格式，目前只支持coco形式
        split_flag: 分割的尺度，list形式
        reverse: True则以一张图片的最小图片作为参考标准，若为False则以最大图片作为参考标准
    Returns:

    """
    split_flag = [0] + split_flag + [inf]
    save_flag = defaultdict(list)
    if types == "coco":
        coco = COCO(path)
        # TODO 注意命名格式，或者可以自己加一个
        filename = path.split("/")[-1].split(".")[0]
        classes, classes_id = id2name(coco)
        print(classes)
        print("class_ids:", classes_id)

        img_ids = coco.getImgIds()
        print(len(img_ids))
        print("loading annotations...")
        for imgId in tqdm(img_ids):
            img = coco.loadImgs(imgId)[0]
            annid = coco.getAnnIds(imgIds=img["id"], iscrowd=None)
            ann = coco.loadAnns(annid)
            scale = 0
            for ann_ in ann:
                height, width = ann_["bbox"][2], ann_["bbox"][3]
                if reverse:
                    scale = min(scale, min(height, width))
                else:
                    scale = max(scale, max(height, width))
            for i in range(1, len(split_flag)):
                # breakpoint()
                if split_flag[i - 1] <= scale < split_flag[i]:
                    save_flag[split_flag[i]].append([img, ann])
    else:
        raise ValueError("types must be cocos ")
    # breakpoint()
    classes_ids = coco.getCatIds(catNms=["head"])
    classes_list = coco.loadCats(classes_ids)
    for key in save_flag.keys():
        if reverse:
            print(
                "saveing {}_min_{}.json, the length is {}".format(
                    filename, key, len(save_flag[key])
                )
            )
            save_path = "data/{}_min_{}.json".format(filename, key)
        else:
            print(
                "saveing {}_max_{}.json, the length is {}".format(
                    filename, key, len(save_flag[key])
                )
            )
            save_path = "data/{}_max_{}.json".format(filename, key)
        _save_json_data(save_path, classes_list, save_flag[key])


def _save_json_data(save_path, classes_list, save_flag):
    coco_sub = dict()
    coco_sub["info"] = dict()
    coco_sub["licenses"] = []
    coco_sub["images"] = []
    coco_sub["type"] = "instances"
    coco_sub["annotations"] = []
    coco_sub["categories"] = []
    # 以下非必须,为coco数据集的前缀信息
    coco_sub["info"]["description"] = "COCO 2017 sub Dataset"
    coco_sub["info"]["url"] = "https://www.cnblogs.com/lhdb/"
    coco_sub["info"]["version"] = "1.0"
    coco_sub["info"]["year"] = 2020
    coco_sub["info"]["contributor"] = "smh"
    coco_sub["info"]["date_created"] = "2020-7-1 10:06"
    sub_license = dict()
    sub_license["url"] = "https://www.cnblogs.com/lhdb/"
    sub_license["id"] = 1
    sub_license["name"] = "Attribution-NonCommercial-ShareAlike License"
    coco_sub["licenses"].append(sub_license)
    # 以下为必须插入信息,包括image、annotations、categories三个字段
    # 插入image信息
    # breakpoint()

    # 插入annotation信息
    for i in range(len(save_flag)):
        coco_sub["images"].append(save_flag[i][0])
        coco_sub["annotations"].extend(save_flag[i][1])
    # 插入categories信息
    coco_sub["categories"].extend(classes_list)
    # 自此所有该插入的数据就已经插入完毕啦٩( ๑╹ ꇴ╹)۶
    # 最后一步，保存数据
    json.dump(coco_sub, open(save_path, "w"))


import time

if __name__ == "__main__":
    a = time.time()
    annFile = "data/phcl/annotations/widerface_val.json"
    # annFile = 'data/phcl/annotations/license-plate_val.json'
    # annFile = ['data/phcl/annotations/widerface_val.json', 'data/phcl/annotations/license-plate_val.json']
    # data = load_dataset(path=annFile, types='coco')
    # scale = data[:,0]*data[:, 1]
    split_dataset(
        path=annFile, types="coco", split_flag=[32, 64, 128, 256], reverse=False
    )
    # ()
    b = time.time()
    print("time consume {}".format(b - a))
